import os
import glob
from PIL import Image
import numpy as np
import random
from typing import Union, Tuple
from pathlib import Path
 
 
data_path = r"/Users/summy/data/garbage_classification" # 替换成自己的数据集路径
# img_list = os.listdir(data_path)
# # random.shuffle(img_list)  # 打乱图片顺序
# result = []
# # for file in img_list[:5000]:  # 根据数据集大小选取部分图片，如果数据集总量没有5000会报错
# for file in img_list:  # 如果数据集太大会很耗内存和时间，可以选择部分图片
#     img = Image.open(os.path.join(data_path, file)).convert('RGB').resize((128, 128)) # (128, 128, 3)
#     img = np.array(img).astype(np.uint8)
#     img = img / 255.
#     result.append(img)
 
# # print(np.shape(result))
# result = np.concatenate(result, axis=-1)
# print(f'result:{result.shape}')
# mean = np.mean(result, axis=(0, 1, 2))   # 对RGB三通道分别求均值
# std = np.std(result, axis=(0, 1, 2))
# print(f'mean:{mean}')
# print(f'std:{std}')


def read_img(img_path, imgsz: Tuple[int, int]) -> np.ndarray:
    """返回的是归一化后的numpy数组"""
    img = Image.open(img_path).convert('RGB').resize(imgsz)  # (h, w, c)
    return np.array(img).astype(np.uint8) / 255.


def get_img_mean_std(data_path: Union[str, Path], imgsz: Union[int, Tuple[int, int]] = (224, 224), 
                     img_prefix: str = 'jpg', sample: Union[int, float] = 2048) -> Tuple[np.ndarray, np.ndarray]:
    """
    计算图片数据集的均值和标准差

    Args:
        data_path: 数据集路径
        img_size: 图片尺寸
        img_prefix: 图片后缀名(jpg, png, jpeg等)
        sample: 样本数量，可以是整数或浮点数，如果为浮点数，则表示样本占比
    """
    if isinstance(data_path, str):
        data_path = Path(data_path)

    if isinstance(imgsz, int):
        imgsz = (imgsz, imgsz)
    
    img_list = list(data_path.glob(f'**/*.{img_prefix}'))
    if isinstance(sample, int):
        if sample < len(img_list):
            img_list = random.sample(img_list, sample)
    elif sample < 1.:
        sample = int(sample * len(img_list))
        img_list = random.sample(img_list, sample)

    result = np.stack(list(map(lambda x: read_img(x, imgsz), img_list)))
    mean = np.mean(result, axis=(0, 1, 2))
    std = np.std(result, axis=(0, 1, 2))
    return mean, std


if __name__ == '__main__':
    mean, std = get_img_mean_std(data_path, imgsz=28, sample=.2)
    print('mean:', mean, 'std:', std)
